import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataloader import MyDataset
from sampler import base_sampler
from deepctr_torch.models import deepfm


class BCE(nn.Module):
    def __init__(self, model, config):
        super().__init__()
        self.model = model
        self.config = config
        self.bceloss = torch.nn.BCEWithLogitsLoss()

    def forward(self, batch_data):
        batch_user, batch_item, label = batch_data
        batch_user = batch_user.to(self.config["device"])
        batch_item = batch_item.to(self.config["device"])
        label = label.to(self.config["device"])
        (users_emb, items_emb, userEmb0, itemmb0) = self.model.get_user_item_embedding(batch_user.long(), batch_item.long())
        reg_loss = (1 / 2) * (userEmb0.norm(2).pow(2) +
                              itemmb0.norm(2).pow(2)) / float(len(batch_user))
        scores = torch.mul(users_emb, items_emb)
        scores = torch.sum(scores, dim=1)
        loss = self.bceloss(scores, label)
        return loss + reg_loss * self.config["weight_decay"]



class BPR_for_user(nn.Module):

    def __init__(self, model, config, dataset):
        super().__init__()
        self.model: deepfm = model
        self.config = config
        self.dataset: MyDataset = dataset
        self.sampler: base_sampler = dataset.sampler
        self.mat = dataset.train_mat.tocsr()

        # self.model:
        # self.dataset
        # self.sampler

    def forward(self, batch_data: torch.int32):
        """
        输入：一个batch的用户
        得到这些user的正样本
        对每个正样本进行负采样  得到每行为 1个正样本 + 多个负样本 的矩阵

        得到正样本的score list，得到负样本的 score mat (有多个负样本)

        :param batch_data:
        :return:
        """
        self.model.train()

        # 取出 batch_user 的交互历史
        batch_uid = batch_data[0]
        batch_mat = self.mat[batch_uid].tocoo()
        indptr = batch_mat.tocsr().indptr

        batch_user_range = batch_mat.row  # 行号会被重置为 range()
        batch_pos = batch_mat.col
        pos_num = batch_mat.sum(1).A.squeeze()  # |P|

        # to gpu
        batch_user_range = torch.LongTensor(batch_user_range)
        batch_pos = torch.from_numpy(batch_pos).reshape(-1, 1)
        pos_num = torch.from_numpy(pos_num).to(self.config["device"])
        indptr = torch.from_numpy(indptr).to(self.config["device"])

        # 按用户进行负采样
        batch_neg_mat, prob_neg_mat = self.sampler(batch_uid.to(self.config["device"]),
                                                   model=self.model)  # batch_neg 是一个矩阵, torch.int32

        # 下面的操作都是按用户分布重复各自重复
        batch_neg_mat = batch_neg_mat.index_select(0, batch_user_range.to(self.config["device"]))
        prob_neg_mat = prob_neg_mat.index_select(0, batch_user_range.to(self.config["device"]))
        pos_num = pos_num.index_select(0, batch_user_range.to(self.config["device"])).reshape(-1, 1)
        batch_user = batch_uid.index_select(0, batch_user_range).reshape(-1, 1)
        #

        batch_neg = batch_neg_mat.reshape(-1, 1).cpu()
        if self.dataset.has_sideinfo:
            """
            torch.tensor -> pandas.DataFrame -> torch.tensor
            """
            batch_user = pd.merge(pd.DataFrame(batch_user.numpy(), columns=["userInt"]), self.dataset.users_info,
                                  how="left", on="userInt")
            batch_pos = pd.merge(pd.DataFrame(batch_pos.numpy(), columns=["newsInt"]), self.dataset.items_info,
                                 how="left", on="newsInt")
            batch_neg = pd.merge(pd.DataFrame(batch_neg.numpy(), columns=["newsInt"]), self.dataset.items_info,
                                 how="left", on="newsInt")
            batch_user = torch.from_numpy(batch_user.to_numpy())
            batch_pos = torch.from_numpy(batch_pos.to_numpy())
            batch_neg = torch.from_numpy(batch_neg.to_numpy())
        batch_user = batch_user.reshape(-1, len(batch_user[0])).to(self.config["device"])
        batch_pos = batch_pos.reshape(-1, len(batch_pos[0])).to(self.config["device"])
        batch_neg = batch_neg.reshape(-1, len(batch_neg[0])).to(self.config["device"])

        # predict
        if self.config["model_name"] == "deepfm":
            # 拼接，统一到torch.tensor
            samples_pos = torch.cat((batch_user, batch_pos), dim=1)
            samples_neg = torch.cat(
                (batch_user.repeat(1, batch_neg_mat.shape[1]).reshape(-1, len(batch_user[0])), batch_neg), dim=1)
            scores_pos = self.model(samples_pos)
            scores_neg = self.model(samples_neg)
            scores_neg_mat = scores_neg.reshape(scores_pos.shape[0], -1)
        elif self.config["model_name"] == "mf":
            scores_pos, scores_neg_mat = self.model(batch_user, batch_pos,
                                                    batch_neg_mat)  # neg_rat是个矩阵   因为一个正样本对应多个负样本
        else:
            raise NotImplementedError

        pred = scores_pos - scores_neg_mat

        weight_loss = torch.negative(F.logsigmoid(pred))
        if self.config["reduction"]:
            loss = torch.sum(weight_loss, dim=-1).mean(-1)
        else:
            loss = torch.sum(weight_loss, dim=-1).sum(-1)

        if self.config["model_name"] == "deepfm":
            reg_loss = self.model.get_regularization_loss()
            return loss + reg_loss * self.config["weight_decay"]
        elif self.config["model_name"] == "mf":
            return loss
        else:
            raise NotImplementedError



class SAUC_for_user(nn.Module):

    def __init__(self, model, config, dataset):
        super().__init__()
        self.model: deepfm = model
        self.config = config
        self.dataset: MyDataset = dataset
        self.sampler: base_sampler = dataset.sampler
        self.mat = dataset.train_mat.tocsr()
        if config["pos_rank_weighted"] is True:
            self._rank_weight_pre = torch.tensor([self.cal_n(x) if x > 1 else 1 for x in range(dataset.items_num + 1)])

    def cal_n(self, n):
        vec = 1 / torch.arange(0, n)
        return vec[1:].sum()

    def forward(self, batch_data: torch.int32):
        """
        输入：一个batch的用户
        得到这些user的正样本
        对每个正样本进行负采样  得到每行为 1个正样本 + 多个负样本 的矩阵

        得到正样本的score list，得到负样本的 score mat (有多个负样本)

        :param batch_data:
        :return:
        """
        self.model.train()

        # 取出 batch_user 的交互历史
        batch_uid = batch_data[0]
        batch_mat = self.mat[batch_uid].tocoo()
        indptr = batch_mat.tocsr().indptr

        batch_user_range = batch_mat.row  # 行号会被重置为 range()
        batch_pos = batch_mat.col
        pos_num = batch_mat.sum(1).A.squeeze()  # |P|

        # to gpu
        batch_user_range = torch.LongTensor(batch_user_range)
        batch_pos = torch.from_numpy(batch_pos).reshape(-1, 1)
        pos_num = torch.from_numpy(pos_num).to(self.config["device"])
        indptr = torch.from_numpy(indptr).to(self.config["device"])

        # 按用户进行负采样
        batch_neg_mat, prob_neg_mat = self.sampler(batch_uid.to(self.config["device"]), model=self.model)  # batch_neg 是一个矩阵, torch.int32

        # 下面的操作都是按用户分布重复各自重复
        batch_neg_mat = batch_neg_mat.index_select(0, batch_user_range.to(self.config["device"]))
        prob_neg_mat = prob_neg_mat.index_select(0, batch_user_range.to(self.config["device"]))
        pos_num = pos_num.index_select(0, batch_user_range.to(self.config["device"])).reshape(-1, 1)
        batch_user = batch_uid.index_select(0, batch_user_range).reshape(-1, 1)


        batch_neg = batch_neg_mat.reshape(-1, 1).cpu()
        if self.dataset.has_sideinfo:
            """
            torch.tensor -> pandas.DataFrame -> torch.tensor
            """
            batch_user = pd.merge(pd.DataFrame(batch_user.numpy(), columns=["userInt"]), self.dataset.users_info,
                                  how="left", on="userInt")
            batch_pos = pd.merge(pd.DataFrame(batch_pos.numpy(), columns=["newsInt"]), self.dataset.items_info,
                                 how="left", on="newsInt")
            batch_neg = pd.merge(pd.DataFrame(batch_neg.numpy(), columns=["newsInt"]), self.dataset.items_info,
                                 how="left", on="newsInt")
            batch_user = torch.from_numpy(batch_user.to_numpy())
            batch_pos = torch.from_numpy(batch_pos.to_numpy())
            batch_neg = torch.from_numpy(batch_neg.to_numpy())
        batch_user = batch_user.reshape(-1, len(batch_user[0])).to(self.config["device"])
        batch_pos = batch_pos.reshape(-1, len(batch_pos[0])).to(self.config["device"])
        batch_neg = batch_neg.reshape(-1, len(batch_neg[0])).to(self.config["device"])


        # predict
        if self.config["model_name"] == "deepfm":
            # 拼接，统一到torch.tensor
            samples_pos = torch.cat((batch_user, batch_pos), dim=1)
            samples_neg = torch.cat(
                (batch_user.repeat(1, batch_neg_mat.shape[1]).reshape(-1, len(batch_user[0])), batch_neg), dim=1)
            scores_pos = self.model(samples_pos)
            scores_neg = self.model(samples_neg)
            scores_neg_mat = scores_neg.reshape(scores_pos.shape[0], -1)
        elif self.config["model_name"] == "mf":
            scores_pos, scores_neg_mat = self.model(batch_user, batch_pos,
                                                    batch_neg_mat)  # neg_rat是个矩阵   因为一个正样本对应多个负样本
        else:
            raise NotImplementedError


        pred = scores_pos - scores_neg_mat

        sauc_weight = 1 / (pos_num * pos_num)

        # importance = importance.to(self.config["device"])
        # pos_rank = pos_rank.to(self.config["device"])
        score_mat = torch.sigmoid(pred / self.config["tau"])
        ui_score = sauc_weight * torch.sum(score_mat, dim=1).unsqueeze(1)

        if self.config["reduction"]:
            loss = 1 - ui_score.squeeze().mean(-1)
        else:
            loss = len(batch_uid) - ui_score.squeeze().sum(-1)

        if self.config["model_name"] == "deepfm":
            reg_loss = self.model.get_regularization_loss()
            return loss + reg_loss * self.config["weight_decay"]
        elif self.config["model_name"] == "mf":
            return loss
        else:
            raise NotImplementedError


def getLoss(lossName, model, config, dataset):
    if lossName == "bce":
        return BCE(model, config)
    elif lossName == "bpr_for_user":
        return BPR_for_user(model, config, dataset)
    elif lossName == "sauc_for_user":
        return SAUC_for_user(model, config, dataset)
    else:
        raise NotImplementedError(f"loss {lossName} haven't been implemented!!!")
